package com.feidee.fdhadoop.yarn;

import com.feidee.fdcommon.util.ToolUtils;
import com.feidee.fdhadoop.kerberos.KerberosAuthentication;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.yarn.api.records.ApplicationReport;
import org.apache.hadoop.yarn.api.records.YarnApplicationState;
import org.apache.hadoop.yarn.client.api.YarnClient;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.apache.hadoop.yarn.util.ConverterUtils;

import java.io.IOException;
import java.util.EnumSet;
import java.util.List;
import java.util.Set;

/**
 * @Author caokaizhi
 * @Date 2018-02-27 15:46
 */
public class YarnUtil {

    static {
        KerberosAuthentication.authenticate();
    }
    private static YarnClient client;

    public static void init(Configuration conf){
        if(client == null){
            client = YarnClient.createYarnClient();
            client.init(conf);
            client.start();
        }
    }

    public static void close() throws IOException {
        if(client != null){
            client.close();
        }
    }

    public static List<ApplicationReport>  getApplications(Set<String> queues, Set<String> users, Set<String> applicationTypes, EnumSet<YarnApplicationState> applicationStates) throws IOException, YarnException {
        List<ApplicationReport> appsReport = null;
        if(client != null) {
            KerberosAuthentication.authenticate();
            appsReport = client.getApplications(queues, users, applicationTypes, applicationStates);
        }
        return appsReport;
    }

    public static ApplicationReport getApplication(Set<String> queues, Set<String> users, Set<String> applicationTypes, EnumSet<YarnApplicationState> applicationStates, String name) throws IOException, YarnException {
        List<ApplicationReport> appsReport = YarnUtil.getApplications(queues,users,applicationTypes,applicationStates);
        if(appsReport != null) {
            for (ApplicationReport appReport : appsReport) {
                if(appReport.getName().equals(name)){
                    return  appReport;
                }
            }
        }
        return null;
    }

    public static ApplicationReport getApplication(String appid) throws IOException, YarnException {
        ApplicationReport appReport = null;
        if(client != null && ToolUtils.isNotNull(appid)) {
            KerberosAuthentication.authenticate();
            appReport = client.getApplicationReport(ConverterUtils.toApplicationId(appid));
        }
        return appReport;
    }


    public static void killApplication(Set<String> queues, Set<String> users, Set<String> applicationTypes, EnumSet<YarnApplicationState> applicationStates, String name) throws YarnException, IOException{
        ApplicationReport appReport = YarnUtil.getApplication(queues,users,applicationTypes,applicationStates,name);
        if(appReport != null && appReport.getYarnApplicationState() != YarnApplicationState.FINISHED
                && appReport.getYarnApplicationState() != YarnApplicationState.KILLED
                && appReport.getYarnApplicationState() != YarnApplicationState.FAILED){
            client.killApplication(appReport.getApplicationId());
        }
    }

    public static void killApplication(String appid) throws YarnException, IOException{
        ApplicationReport  appReport = getApplication(appid);
        if(appReport != null && appReport.getYarnApplicationState() != YarnApplicationState.FINISHED
                && appReport.getYarnApplicationState() != YarnApplicationState.KILLED
                && appReport.getYarnApplicationState() != YarnApplicationState.FAILED){
            client.killApplication(appReport.getApplicationId());
        }
    }
}
